package com.example.leetcode.graph.common.unionfind.onedimension.pro.optimize;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 将「路径压缩优化」和「按秩合并优化」合并后形成的「并查集」的实现方式
 */
public class UnionFind {
    int[] root;
    // 添加了 rank 数组来记录每个顶点的高度，也就是每个顶点的「秩」
    int[] rank;
    // 连通分量个数
    int count;

    public UnionFind(int size) {
        root = new int[size];
        rank = new int[size];
        // 开始时连通分量数量即为节点数量
        count = size;
        for (int i = 0; i < size; i++) {
            // 初始时父结点为自身
            root[i] = i;
            // 树高为1
            rank[i] = 1; // 一开始每个顶点的初始「秩」为1，因为它们只有自己本身的一个顶点。
        }
    }

    /**
     * 获取图的连通分量数量
     * 连通分量数量即为根节点的数量，即父结点为自身的节点数量
     *
     * @return
     */
    public int rootCount() {
        int count = 0;
        for (int i = 0; i < root.length; i++) {
            if (root[i] == i) {
                count++;
            }
        }
        return count;
    }

    /**
     * 直接获取连通分量个数
     *
     * @return
     */
    public int getCount() {
        return count;
    }

    // 此处的 find 函数与路径压优化缩版本的 find 函数一样。
    public int find(int x) {
        // 如果当前节点是根节点，返回当前节点
        if (x == root[x]) {
            return x;
        }
        // 简便写法 return root[x] = find(root[x]);
        // 找到根节点
        // 递归查找根节点
        int result = find(root[x]);
        // 将当前节点的父节点替换为根节点，进行路径压缩
        root[x] = result;
        // 返回根节点
        return result;
    }

    // 按秩合并优化的 union 函数
    public void union(int x, int y) {
        int rootX = find(x);
        int rootY = find(y);
        if (rootX != rootY) {
            // 将秩小的分量合并到秩大的分量下，避免形成链表
            if (rank[rootX] > rank[rootY]) {
                root[rootY] = rootX;
            } else if (rank[rootX] < rank[rootY]) {
                root[rootX] = rootY;
            } else {
                root[rootY] = rootX;
                rank[rootX] += 1;
            }
            // 每合并两个连通分量后，连通分量少一
            count--;
        }
    }

    /**
     * 判断两个元素是否位于同一个连通分量中
     *
     * @param x
     * @param y
     * @return
     */
    public boolean isConnected(int x, int y) {
        return find(x) == find(y);
    }

    /**
     * 获取连通分量的根元素
     *
     * @return
     */
    public List<Integer> getAllRoot() {
        List<Integer> result = new ArrayList<>();
        for (int i = 0; i < root.length; i++) {
            if (root[i] == i) {
                result.add(i);
            }
        }
        return result;
    }

    /**
     * 获取所有的连通分量
     *
     * @return
     */
    public Map<Integer, List<Integer>> getAllConnectedComponent() {
        Map<Integer, List<Integer>> map = new HashMap<>();
        for (int i = 0; i < root.length; i++) {
            int root = find(i);
            List<Integer> component;
            component = map.getOrDefault(root, new ArrayList<>());
            component.add(i);
            map.put(root, component);
        }
        return map;
    }

}